[None][fix] Fix one-model MTP KV cache accounting - #17264

Can we make this boundary decision group-uniform under attention DP? Each rank schedules different requests, so one rank can set runtime_draft_len=0 and return through skip_drafting() while a peer keeps K and enters DeepseekV3MTPHead's LM-head-TP all-gather. That leaves an unmatched collective and hangs the group. Please OR the unsafe flag across the TP/ADP group before changing the batch-wide draft length, and add a two-rank near-limit test (optional).

Good catch!

_handle_dynamic_draft_len() now computes a local runtime_draft_len first, ORs the unsafe flag across the TP/ADP group, and only then mutates runtime_draft_len and request state. All three call sites are already inside if can_queue, which is itself a TP allgather under ADP, so every rank enters the new collective exactly once in the same position; the non-ADP path stays collective-free.

Note that as a result, the diff is quite large: the speculation_permanently_disabled is an early return, and the pad / truncate loop moved to the tail behind if not dynamic_draft_len_enabled: return, and _one_model_mtp_batch_needs_zero_draft() now takes runtime_draft_len as a parameter since it can no longer read the not-yet-committed engine field. Note one deliberate semantic flip inside it: runtime_draft_len == 0 now returns True instead of False, so a locally selected zero also participates in the OR.

One more thing worth pointing at in the diff, since it is easy to read as unrelated: the clear is now if needs_zero_draft or runtime_draft_len == 0.

My first pass at this restructure folded the old unconditional speculation_permanently_disabled clear into the needs_zero_draft block, which is gated on is_mtp_eagle_one_model(). That silently dropped the clear for vanilla MTP and Eagle3 one-model, leaving K python draft tokens alongside runtime_draft_len == 0. Clearing on any committed zero restores it; test_permanent_disable_clears_non_mtp_eagle_draft_tokens is the corresponding test.

This zero-draft decision is lost for dynamic-tree MTP: overlap still supplies next_draft_tokens_device, so _prepare_tp_inputs() classifies the request as extend and then unconditionally resets runtime_draft_len to max_total_draft_tokens for non-linear trees. A near-limit max_draft_len=3, topK=4 request therefore executes all 12 tree tokens after this code selected zero. Can we preserve the explicit zero through TP-input preparation and extend the regression through that stage?

Agreed.

_prepare_tp_inputs() no longer rewrites runtime_draft_len unconditionally for non-linear trees; the normalization is now skipped when the executor already selected zero.

This also explains the comment-only change at can_promote_spec_decode: its old justification no longer holds.

Regression extended past _handle_dynamic_draft_len() and through _prepare_tp_inputs() with a non-None next_draft_tokens_device, and is now tested in test_dynamic_tree_prepare_preserves_explicit_zero_draft.

Can we preserve the empty Python-side draft signal until _handle_dynamic_draft_len() consumes it? After a K=0 iteration, SpecSampler writes py_draft_tokens=[]; filling it here before the handler makes py_needs_onehot_draft_probs stay false. Since that slot produced no proposal row at K=0, the next rejection step gathers stale draft_probs and biases sampling. The scheduler only needs request.draft_tokens, so the Python placeholder can be deferred; please add a rejection-enabled K->0->K regression.

Good catch!

Note that deferring the python placeholder would regress this PR's capacity fix, because get_draft_token_length() reads py_draft_tokens and SchedulerV2, KV allocation, and resource management all budget from it.

Instead, the zero-proposal fact is captured before the placeholder is inserte, and _prepare_tp_inputs() consumes it exactly once. Timing works out because a K=0 request has empty py_draft_tokens and therefore lands in generation_requests, not extend_requests. The flag isn't consumed until the resumed K>0 step, which is when the one-hot row is actually needed.

Covered by test_one_model_mtp_preserves_zero_proposal_signal_for_rejection.

The page reservation can stay intact, but paged prefill still derives its causal KV length from this unchanged indptr/last-page pair (and cached-MLA context does likewise). With cached=31, q_len=4, and offset=-3, appends move to 28..31 while prefill still aligns Q to a 35-token KV, admitting future/stale slots. Can we publish the corrected logical context lengths to the prefill/MLA-context runtime, or reject this extend-context configuration? The added metadata-only SA test never reaches the real extend-context plan.

Agreed, and I took the narrow option: overlap is now disabled for FlashInfer + any speculation mode that uses extend_ctx, rather than publishing corrected logical context lengths.

extend_ctx() returns not issubclass(backend, TrtllmAttention) unless use_one_engine(), so for FlashInfer the gate condition not use_one_engine() is exactly "this run would take the extend-context path" -> one-engine FlashInfer speculation keeps overlap.

A real device-side logical-length contract for paged prefill and cached MLA context (page-boundary rewinds, CUDA-graph-stable buffers, validated against an actual plan/run rather than metadata mutation) is materially larger, so I think it should be punted.

One qualification on severity, since it changes how you may want to weigh this: I could not find a public configuration that reaches the FlashInfer extend-context path today:

extend_ctx requires not use_one_engine(), and combined with support_overlap_scheduler(), that leaves exactly has_draft_model() - EAGLE3 two-model, MTP-Eagle two-model, and DraftTarget two-model. All three are currently unreachable. So I believe the stale-length bug you identified is real but currently latent, and this gate is a guard for retained legacy / internal states rather than a fix for a live path. The test exercises it by mutating that private flag and is labeled as such.

For diff-reading: this is why the inline overlap check in create_py_executor was extracted into _disable_unsupported_speculative_overlap_scheduler().

This fixes only the first overlap generation for metadata without dynamic KV lengths. On every later step the host cached count is still the speculative upper bound, while the accepted-length correction is device-side; VanillaAttentionMetadata has neither kv_lens_cuda nor apply_spec_decode_kv_lens_offsets(). For P=5, width=4, and one accepted token, Vanilla uses cached=9 although the committed boundary is 6, even as positions rewind to 6. Can we implement the dynamic-length contract for these backends or reject overlap speculation for them?

Another good catch!

Overlap speculation now disabled for the vanilla backend, and I removed the py_decoding_iter == 0 special case outright so nothing implies later steps are safe.

This is covered by the new TestUnsupportedSpeculativeOverlapScheduler, which asserts vanilla + speculation disables overlap, and that an already-disabled config is left alone.

Can we avoid testing spec_dec_mode here before AutoDeploy has resolved the checkpoint's MTP layer count? num_nextn_predict_layers is init=False; the input validator remaps it to max_draft_len, leaving it None, and this property treats None as one layer. AD also skips ModelLoader, the only pre-executor writer of that field. As a result every allowed AD MTP config—including the registered smoke/accuracy tests—hits this NotImplementedError. Please resolve the actual mode before this predicate or preserve the existing AD MTP path.

Done.